from fastapi import APIRouter, Depends, status, HTTPException
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from sqlalchemy.orm import Session
from datetime import datetime, timedelta
from typing import Optional
from jose import jwt
import bcrypt
import os
import secrets
from api.database import get_db
from api.model.user import User
from api.schemas.auth import Token, TokenData, UserCreate, UserResponse, UserLogin, VerifyPasswordRequest, UnifiedVerifyRequest, UnifiedVerifyResponse
from api.utils.response import success, fail
from api.utils.async_helper import run_in_executor
import pyotp
import json
import base64
import logging

# 配置日志
logger = logging.getLogger(__name__)

router = APIRouter()

# JWT配置 - 从环境变量读取
SECRET_KEY = os.getenv("SECRET_KEY", "your-super-secret-key-change-this-in-production")
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="auth/login")

def verify_password(plain_password: str, hashed_password: str) -> bool:
    return bcrypt.checkpw(plain_password.encode('utf-8'), hashed_password.encode('utf-8'))

def get_password_hash(password: str) -> str:
    return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')

def authenticate_user(db: Session, username: str, password: str):
    user = db.query(User).filter(User.username == username).first()
    if not user or not verify_password(password, user.hashed_password):
        return None
    return user

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

async def get_current_user(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="无法验证凭据",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(username=username)
    except jwt.JWTError:
        raise credentials_exception
    user = db.query(User).filter(User.username == token_data.username).first()
    if user is None:
        raise credentials_exception
    return user

def get_current_user_from_token(token: str, db: Session):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            return None
        token_data = TokenData(username=username)
    except jwt.JWTError:
        return None
    user = db.query(User).filter(User.username == token_data.username).first()
    return user

@router.post("/register")
async def register(user: UserCreate, db: Session = Depends(get_db)):
    try:
        db_user = db.query(User).filter(User.username == user.username).first()
        if db_user:
            return fail("用户名已存在")
        
        is_first_user = db.query(User).count() == 0
        # 异步执行密码哈希（CPU密集操作）
        hashed_password = await run_in_executor(get_password_hash, user.password)
        db_user = User(
            username=user.username,
            email=f"{user.username}@example.com",
            hashed_password=hashed_password,
            is_admin=is_first_user
        )
        db.add(db_user)
        db.commit()
        db.refresh(db_user)
        return success(UserResponse.from_orm(db_user), msg="注册成功")
    except Exception as e:
        db.rollback()
        return fail(f"注册失败: {str(e)}")

@router.post("/login")
async def login(form_data: UserLogin, db: Session = Depends(get_db)):
    try:
        # 异步执行认证（包含密码验证）
        user = await run_in_executor(authenticate_user, db, form_data.username, form_data.password)
        if not user:
            return fail("用户名或密码错误")
        
        # 检查账号是否被禁用或删除
        if not user.is_active:
            return fail("账号已被禁用")
        
        if user.delete_time is not None:
            return fail("账号已被删除")

        if user.is_mfa_enabled:
            if not form_data.mfa_token:
                # 当需要MFA但未提供时，返回自定义错误信息
                return fail("需要MFA令牌", data={"mfa_required": True})
            
            totp = pyotp.TOTP(user.mfa_secret)
            if not totp.verify(form_data.mfa_token):
                return fail("无效的MFA令牌")

        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
        access_token = create_access_token(
            data={"sub": user.username}, expires_delta=access_token_expires
        )
        token_data = {"access_token": access_token, "token_type": "bearer"}
        return success(token_data, msg="登录成功")
    except Exception as e:
        return fail(f"登录失败: {str(e)}")

@router.post("/logout")
async def logout():
    return success(msg="已成功退出登录")

@router.get("/me")
async def read_users_me(current_user: User = Depends(get_current_user)):
    try:
        user_data = UserResponse.from_orm(current_user)
        return success(user_data)
    except Exception as e:
        return fail(f"获取用户信息失败: {str(e)}")

@router.post("/verify-password")
async def verify_current_password(
    request: VerifyPasswordRequest,
    current_user: User = Depends(get_current_user)
):
    try:
        # 异步执行密码验证
        is_valid = await run_in_executor(verify_password, request.password, current_user.hashed_password)
        if not is_valid:
            return fail("密码不正确")
        return success(msg="密码验证成功")
    except Exception as e:
        return fail(f"密码验证失败: {str(e)}")

# 临时存储验证令牌（生产环境应使用Redis）
_verification_tokens = {}

# 临时存储已使用的MFA代码（防止重放攻击）
_used_mfa_tokens = {}

def create_verification_token(user_id: int) -> tuple[str, datetime]:
    """创建短期验证令牌（5分钟有效）"""
    token = secrets.token_urlsafe(32)
    expires_at = datetime.utcnow() + timedelta(minutes=5)
    _verification_tokens[token] = {
        "user_id": user_id,
        "expires_at": expires_at
    }
    return token, expires_at

def mark_mfa_token_used(user_id: int, mfa_token: str):
    """标记MFA代码已使用（90秒内不能重复使用）"""
    key = f"{user_id}:{mfa_token}"
    expires_at = datetime.utcnow() + timedelta(seconds=90)
    _used_mfa_tokens[key] = expires_at
    
    # 清理过期的记录
    now = datetime.utcnow()
    expired_keys = [k for k, v in _used_mfa_tokens.items() if now > v]
    for k in expired_keys:
        del _used_mfa_tokens[k]

def is_mfa_token_used(user_id: int, mfa_token: str) -> bool:
    """检查MFA代码是否已被使用"""
    key = f"{user_id}:{mfa_token}"
    if key not in _used_mfa_tokens:
        return False
    
    # 检查是否过期
    if datetime.utcnow() > _used_mfa_tokens[key]:
        del _used_mfa_tokens[key]
        return False
    
    return True

def verify_verification_token(token: str, user_id: int) -> bool:
    """验证令牌是否有效"""
    if token not in _verification_tokens:
        return False
    
    token_data = _verification_tokens[token]
    
    # 检查是否过期
    if datetime.utcnow() > token_data["expires_at"]:
        del _verification_tokens[token]
        return False
    
    # 检查用户ID是否匹配
    if token_data["user_id"] != user_id:
        return False
    
    return True

def invalidate_verification_token(token: str):
    """使验证令牌失效"""
    if token in _verification_tokens:
        del _verification_tokens[token]

@router.post("/verify-unified")
async def verify_unified(
    request: UnifiedVerifyRequest,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """
    统一验证接口，支持密码/MFA/WebAuthn三种验证方式
    返回短期验证令牌用于后续敏感操作
    """
    try:
        verified = False
        
        if request.method == "password":
            # 密码验证
            if not request.password:
                return fail("请提供密码")
            # 异步执行密码验证
            is_valid = await run_in_executor(verify_password, request.password, current_user.hashed_password)
            if is_valid:
                verified = True
            else:
                return fail("密码不正确")
        
        elif request.method == "mfa":
            # MFA验证
            if not current_user.is_mfa_enabled:
                return fail("MFA未启用")
            if not request.mfa_token:
                return fail("请提供MFA代码")
            
            # 检查MFA代码是否已被使用
            if is_mfa_token_used(current_user.id, request.mfa_token):
                logger.warning(f"用户 {current_user.id} 尝试重复使用MFA代码: {request.mfa_token}")
                return fail("该MFA代码已被使用，请等待新代码生成")
            
            totp = pyotp.TOTP(current_user.mfa_secret)
            
            # 使用更宽松的验证窗口（valid_window=2 表示前后各2个30秒窗口，共150秒）
            # 注意：这会降低安全性，建议同步服务器时间后改回valid_window=1
            if totp.verify(request.mfa_token, valid_window=2):
                # 标记为已使用
                mark_mfa_token_used(current_user.id, request.mfa_token)
                verified = True
            else:
                logger.warning(f"MFA验证失败 - 用户ID: {current_user.id}, 输入代码: {request.mfa_token}")
                return fail("MFA代码不正确")
        
        elif request.method == "webauthn":
            # WebAuthn验证
            if not current_user.is_webauthn_enabled:
                return fail("硬件密钥未启用")
            if not request.webauthn_response:
                return fail("请提供硬件密钥验证数据")
            
            # 从webauthn_response中提取challenge验证
            # 这里简化处理，实际应该完整验证签名
            credential_data = request.webauthn_response
            credential_id = credential_data.get("id")
            
            if credential_id == current_user.webauthn_credential_id:
                verified = True
            else:
                return fail("硬件密钥验证失败")
        
        else:
            return fail("不支持的验证方法")
        
        if verified:
            # 生成验证令牌
            token, expires_at = create_verification_token(current_user.id)
            return success({
                "verification_token": token,
                "expires_at": expires_at.isoformat()
            }, msg="验证成功")
        else:
            return fail("验证失败")
    
    except Exception as e:
        return fail(f"验证失败: {str(e)}")
